package com.example.repurchase.prediction;

import org.apache.flink.api.common.eventtime.SerializableTimestampAssigner;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.AggregateFunction;
import org.apache.flink.api.common.functions.RichFlatMapFunction;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.tuple.Tuple3;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.source.SourceFunction;
import org.apache.flink.streaming.api.functions.windowing.ProcessWindowFunction;
import org.apache.flink.streaming.api.windowing.assigners.SlidingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.util.Collector;

import java.time.Duration;
import java.util.Random;

public class UserRepurchasePrediction {
    public static void main(String[] args) throws Exception {
        // 设置执行环境
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(4);
        // 生成模拟购买事件流
        DataStream<PurchaseEvent> purchaseEvents = env.addSource(new PurchaseEventSource())
                .assignTimestampsAndWatermarks(WatermarkStrategy
                        .<PurchaseEvent>forBoundedOutOfOrderness(Duration.ofSeconds(5))
                        .withTimestampAssigner((SerializableTimestampAssigner<PurchaseEvent>)
                                (event, timestamp) -> event.getTimestamp()));
        // 计算每个用户的购买频率
        DataStream<Tuple3<String, Long, Double>> purchaseFrequency = purchaseEvents
                .keyBy(PurchaseEvent::getUserId)
                .window(SlidingEventTimeWindows.of(Time.days(30), Time.days(7)))
                .aggregate(new PurchaseFrequencyAggregate(), new PurchaseFrequencyWindowFunction());
        // 计算每个用户的平均购买金额
        DataStream<Tuple3<String, Long, Double>> avgPurchaseAmount = purchaseEvents
                .keyBy(PurchaseEvent::getUserId)
                .window(SlidingEventTimeWindows.of(Time.days(30), Time.days(7)))
                .aggregate(new AvgPurchaseAmountAggregate(), new AvgPurchaseAmountWindowFunction());
        // 连接两个流并计算复购概率
        DataStream<UserRepurchaseScore> repurchaseScores = purchaseFrequency
                .keyBy(t -> t.f0)
                .connect(avgPurchaseAmount.keyBy(t -> t.f0))
                .flatMap(new CalculateRepurchaseScore());
        // 输出结果
        repurchaseScores.print();
        // 执行作业
        env.execute("User Repurchase Prediction");
    }
    // 购买事件类
    public static class PurchaseEvent {
        private String userId;
        private double amount;
        private long timestamp;

        public PurchaseEvent(String userId, double amount, long timestamp) {
            this.userId = userId;
            this.amount = amount;
            this.timestamp = timestamp;
        }

        public String getUserId() {
            return userId;
        }

        public double getAmount() {
            return amount;
        }

        public long getTimestamp() {
            return timestamp;
        }

        @Override
        public String toString() {
            return "PurchaseEvent{" +
                    "userId='" + userId + '\'' +
                    ", amount=" + amount +
                    ", timestamp=" + timestamp +
                    '}';
        }
    }

    // 购买频率聚合函数
    public static class PurchaseFrequencyAggregate implements AggregateFunction<PurchaseEvent, Long, Long> {
        @Override
        public Long createAccumulator() {
            return 0L;
        }

        @Override
        public Long add(PurchaseEvent value, Long accumulator) {
            return accumulator + 1;
        }

        @Override
        public Long getResult(Long accumulator) {
            return accumulator;
        }

        @Override
        public Long merge(Long a, Long b) {
            return a + b;
        }
    }

    // 购买频率窗口函数
    public static class PurchaseFrequencyWindowFunction extends ProcessWindowFunction<Long, Tuple3<String, Long, Double>, String, TimeWindow> {
        @Override
        public void process(String userId, Context context, Iterable<Long> elements, Collector<Tuple3<String, Long, Double>> out) {
            Long count = elements.iterator().next();
            long windowEnd = context.window().getEnd();
            out.collect(Tuple3.of(userId, windowEnd, count.doubleValue()));
        }
    }

    // 平均购买金额聚合函数
    public static class AvgPurchaseAmountAggregate implements AggregateFunction<PurchaseEvent, Tuple2<Double, Integer>, Double> {
        @Override
        public Tuple2<Double, Integer> createAccumulator() {
            return Tuple2.of(0.0, 0);
        }

        @Override
        public Tuple2<Double, Integer> add(PurchaseEvent value, Tuple2<Double, Integer> accumulator) {
            return Tuple2.of(accumulator.f0 + value.getAmount(), accumulator.f1 + 1);
        }

        @Override
        public Double getResult(Tuple2<Double, Integer> accumulator) {
            return accumulator.f1 == 0 ? 0.0 : accumulator.f0 / accumulator.f1;
        }

        @Override
        public Tuple2<Double, Integer> merge(Tuple2<Double, Integer> a, Tuple2<Double, Integer> b) {
            return Tuple2.of(a.f0 + b.f0, a.f1 + b.f1);
        }
    }

    // 平均购买金额窗口函数
    public static class AvgPurchaseAmountWindowFunction extends ProcessWindowFunction<Double, Tuple3<String, Long, Double>, String, TimeWindow> {
        @Override
        public void process(String userId, Context context, Iterable<Double> elements, Collector<Tuple3<String, Long, Double>> out) {
            Double avgAmount = elements.iterator().next();
            long windowEnd = context.window().getEnd();
            out.collect(Tuple3.of(userId, windowEnd, avgAmount));
        }
    }

    // 计算复购分数
    public static class CalculateRepurchaseScore extends RichFlatMapFunction<Tuple3<String, Long, Double>, UserRepurchaseScore> {
        private transient ValueState<Tuple2<Double, Double>> state;

        @Override
        public void open(Configuration parameters) {
            state = getRuntimeContext().getState(
                    new ValueStateDescriptor<>("repurchase-features", Types.TUPLE(Types.DOUBLE, Types.DOUBLE))
            );
        }

        @Override
        public void flatMap(Tuple3<String, Long, Double> input, Collector<UserRepurchaseScore> out) throws Exception {
            String userId = input.f0;
            long windowEnd = input.f1;
            double value = input.f2;

            Tuple2<Double, Double> features = state.value();
            if (features == null) {
                features = Tuple2.of(null, null);
            }

            // 根据输入流类型更新对应的特征值
            if (getRuntimeContext().getIndexOfThisSubtask() % 2 == 0) {
                features.f0 = value; // 更新频率
            } else {
                features.f1 = value; // 更新金额
            }

            state.update(features);

            // 当两个特征都有值时计算复购分数
            if (features.f0 != null && features.f1 != null) {
                double score = calculateRepurchaseScore(features.f0, features.f1);
                out.collect(new UserRepurchaseScore(userId, windowEnd, features.f0, features.f1, score));
            }
        }

        private double calculateRepurchaseScore(double frequency, double avgAmount) {
            // 简单的复购分数计算逻辑，可以替换为更复杂的机器学习模型
            double baseScore = 0.3;
            double frequencyFactor = Math.min(frequency / 10, 0.5);
            double amountFactor = Math.min(avgAmount / 1000, 0.2);
            return baseScore + frequencyFactor + amountFactor;
        }
    }

    // 用户复购分数类
    public static class UserRepurchaseScore {
        private String userId;
        private long windowEndTime;
        private double purchaseFrequency;
        private double avgPurchaseAmount;
        private double repurchaseScore;

        public UserRepurchaseScore() {}

        public UserRepurchaseScore(String userId, long windowEndTime, double purchaseFrequency, double avgPurchaseAmount, double repurchaseScore) {
            this.userId = userId;
            this.windowEndTime = windowEndTime;
            this.purchaseFrequency = purchaseFrequency;
            this.avgPurchaseAmount = avgPurchaseAmount;
            this.repurchaseScore = repurchaseScore;
        }

        @Override
        public String toString() {
            return "UserRepurchaseScore{" +
                    "userId='" + userId + '\'' +
                    ", windowEndTime=" + windowEndTime +
                    ", purchaseFrequency=" + purchaseFrequency +
                    ", avgPurchaseAmount=" + avgPurchaseAmount +
                    ", repurchaseScore=" + String.format("%.2f", repurchaseScore) +
                    '}';
        }
    }

    // 模拟购买事件源
    public static class PurchaseEventSource implements SourceFunction<PurchaseEvent> {
        private boolean running = true;
        private final Random random = new Random();
        private final String[] userIds = {"user1", "user2", "user3", "user4", "user5", "user6", "user7", "user8", "user9", "user10"};

        @Override
        public void run(SourceContext<PurchaseEvent> ctx) throws Exception {
            while (running) {
                String userId = userIds[random.nextInt(userIds.length)];
                double amount = random.nextDouble() * 1000;
                long timestamp = System.currentTimeMillis();

                ctx.collect(new PurchaseEvent(userId, amount, timestamp));

                // 控制数据生成速率
                Thread.sleep(100);
            }
        }

        @Override
        public void cancel() {
            running = false;
        }
    }
}    